# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: donggu

# std
import collections
import struct
import os
import gzip

# 3rd party
import tqdm

# internal
from coin.flow.topic_record import TopicRecordWriter
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog, OrderEvent
import coin.proto.coin_order_enums_pb2 as coin_order
from coin.proto.support.coin_order_stateless_pb2 import StatelessOrderEvent, OrderStateTrace

from absl import flags, app

FLAGS = flags.FLAGS
flags.DEFINE_string('date', None, 'example: --date 20190708')
flags.mark_flag_as_required('date')
flags.register_validator('date', lambda value: value.isdigit(), message='date example: 20190708')

SRC_ROOT = '/remote/iosg/strat-1/buckets/log.raw.coin/live/og_proto_log/'
DST_ROOT = '/remote/iosg/coin/data/og_proto_stateless_log/'
MAGIC = b'\x00\x31\x41\x59\x00\x01\x00\x00'


def read_orders_from_file(in_path):
  if in_path.endswith('.gz'):
    open_fn = gzip.open
  else:
    open_fn = open

  with open_fn(in_path, "rb") as f:
    magic = f.read(8)
    if magic != MAGIC:
      # printing within tqdm, requires \n for clear print
      raise RuntimeError('magic dismatch')

    while True:
      length = f.read(8)
      if length == b'':
        return

      length = struct.unpack('<Q', length)[0]
      f.read(4 + 32)  # unused metadata
      data = f.read(length - 32)
      f.read(4)  # unused crc

      e = OrderGatewayLog()
      e.ParseFromString(data)
      yield e


def convert_file(in_path, out_path):
  trace = collections.defaultdict(lambda: OrderStateTrace())

  with TopicRecordWriter(out_path) as f_out:
    for order in read_orders_from_file(in_path):
      if order.type != OrderGatewayLog.ORDER_EVENT:
        continue

      ev = order.event

      if ev.event_time == 0:
        print('order:', order)
        raise RuntimeError('event_time is 0')

      # order_state
      if ev.order_state == coin_order.CREATED_ORDER:
        if trace[ev.internal_order_id].created_time == 0:
          trace[ev.internal_order_id].created_time = ev.event_time
      elif ev.order_state == coin_order.PENDING_ORDER:
        if trace[ev.internal_order_id].pending_time == 0:
          trace[ev.internal_order_id].pending_time = ev.event_time
      elif ev.order_state == coin_order.WORKING_ORDER:
        if trace[ev.internal_order_id].working_time == 0:
          trace[ev.internal_order_id].working_time = ev.event_time
      elif ev.order_state == coin_order.DEAD_ORDER:
        if trace[ev.internal_order_id].dead_time == 0:
          trace[ev.internal_order_id].dead_time = ev.event_time

      # order_event for order_state
      if ev.type == OrderEvent.ORDER_SUBMITTED:
        if trace[ev.internal_order_id].pending_time == 0:
          trace[ev.internal_order_id].pending_time = ev.event_time
      elif ev.type == OrderEvent.ORDER_ACCEPTED:
        if trace[ev.internal_order_id].working_time == 0:
          trace[ev.internal_order_id].working_time = ev.event_time
      elif ev.type in [
          OrderEvent.ORDER_EXPIRED,
          OrderEvent.ORDER_REJECTED,
          OrderEvent.CANCEL_CONFIRMED,
          OrderEvent.ORDER_AUTO_CANCELED
      ]:
        if trace[ev.internal_order_id].dead_time == 0:
          trace[ev.internal_order_id].dead_time = ev.event_time

      # order_event
      if ev.type == OrderEvent.UNKNOWN_ORDER_EVENT:
        if trace[ev.internal_order_id].unknown_order_event_time == 0:
          trace[ev.internal_order_id].unknown_order_event_time = ev.event_time
      elif ev.type == OrderEvent.ORDER_SUBMITTED:
        if trace[ev.internal_order_id].order_submitted_time == 0:
          trace[ev.internal_order_id].order_submitted_time = ev.event_time
      elif ev.type == OrderEvent.ORDER_ACCEPTED:
        if trace[ev.internal_order_id].order_accepted_time == 0:
          trace[ev.internal_order_id].order_accepted_time = ev.event_time
      elif ev.type == OrderEvent.ORDER_ERROR:
        if trace[ev.internal_order_id].order_error_time == 0:
          trace[ev.internal_order_id].order_error_time = ev.event_time
      elif ev.type == OrderEvent.ORDER_FILLED:  # special case
        lst = trace[ev.internal_order_id].order_filled.add()
        lst.timestamp = ev.event_time
        lst.qty = ev.fill_qty
        lst.price = ev.fill_price
      elif ev.type == OrderEvent.ORDER_AUTO_CANCELED:
        if trace[ev.internal_order_id].order_auto_canceled_time == 0:
          trace[ev.internal_order_id].order_auto_canceled_time = ev.event_time
      elif ev.type == OrderEvent.ORDER_AMENDED:
        if trace[ev.internal_order_id].order_amended_time == 0:
          trace[ev.internal_order_id].order_amended_time = ev.event_time
      elif ev.type == OrderEvent.ORDER_AMEND_ACCEPTED:
        if trace[ev.internal_order_id].order_amend_accepted_time == 0:
          trace[ev.internal_order_id].order_amend_accepted_time = ev.event_time
      elif ev.type == OrderEvent.ORDER_REJECTED:
        if trace[ev.internal_order_id].order_rejected_time == 0:
          trace[ev.internal_order_id].order_rejected_time = ev.event_time
      elif ev.type == OrderEvent.ORDER_EXPIRED:
        if trace[ev.internal_order_id].order_expired_time == 0:
          trace[ev.internal_order_id].order_expired_time = ev.event_time
      elif ev.type == OrderEvent.CANCEL_SUBMITTED:
        if trace[ev.internal_order_id].cancel_submitted_time == 0:
          trace[ev.internal_order_id].cancel_submitted_time = ev.event_time
      elif ev.type == OrderEvent.CANCEL_ACCEPTED:
        if trace[ev.internal_order_id].cancel_accepted_time == 0:
          trace[ev.internal_order_id].cancel_accepted_time = ev.event_time
      elif ev.type == OrderEvent.CANCEL_ERROR:
        if trace[ev.internal_order_id].cancel_error_time == 0:
          trace[ev.internal_order_id].cancel_error_time = ev.event_time
      elif ev.type == OrderEvent.CANCEL_CONFIRMED:
        if trace[ev.internal_order_id].cancel_confirmed_time == 0:
          trace[ev.internal_order_id].cancel_confirmed_time = ev.event_time

      out = StatelessOrderEvent()
      out.event.CopyFrom(ev)
      out.state.CopyFrom(trace[ev.internal_order_id])

      f_out.write(None, 0, None, ev.event_time, out.SerializeToString())


def get_paths(date):
  # return [(in_path, out_path), ..]
  strategies = os.listdir(SRC_ROOT)
  paths = []

  for strat in tqdm.tqdm(strategies, desc='ls strategy'):
    try:
      dir_ = os.path.join(SRC_ROOT, strat, date)
      paths_ = [(os.path.join(SRC_ROOT, strat, date, filename),
                 os.path.join(DST_ROOT, strat, date, filename)) for filename in os.listdir(dir_)]

      paths.extend(paths_)
    except FileNotFoundError:
      # data does not exist
      pass

  return paths


def main(argv):
  paths = get_paths(FLAGS.date)

  for (in_path, out_path) in tqdm.tqdm(paths, desc='convert'):
    # abc.gz => abc
    if out_path.endswith('.gz'):
      out_path = out_path[:-3]

    out_dir = os.path.dirname(out_path)
    os.makedirs(out_dir, exist_ok=True)

    try:
      convert_file(in_path, out_path)
    except FileExistsError:
      pass
    except Exception as err:
      print('error from file: ', in_path)
      print(err)


if __name__ == '__main__':
  app.run(main)
